// SPDX-FileCopyrightText: 2025 UnionTech Software Technology Co., Ltd.
// SPDX-License-Identifier: MIT

package ratelimit

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
)

// testWithTimeout 为时间敏感的测试提供超时保护
func testWithTimeout(t *testing.T, timeout time.Duration, testFunc func()) {
	done := make(chan struct{})
	go func() {
		defer close(done)
		testFunc()
	}()

	select {
	case <-done:
		// 测试完成
	case <-time.After(timeout):
		t.Fatal("Test timed out")
	}
}

func TestNew(t *testing.T) {
	rate := 10
	limiter := New(rate)
	if limiter == nil {
		t.Fatal("Expected limiter to be created, got nil")
	}

	if limiter.rate != rate {
		t.Errorf("Expected rate %d, got %d", rate, limiter.rate)
	}

	if limiter.tokens != rate {
		t.Errorf("Expected initial tokens %d, got %d", rate, limiter.tokens)
	}

	if limiter.lastTime.IsZero() {
		t.Error("Expected lastTime to be set")
	}
}

func TestNewWithZeroRate(t *testing.T) {
	limiter := New(0)
	if limiter == nil {
		t.Fatal("Expected limiter to be created with zero rate, got nil")
	}

	if limiter.rate != 0 {
		t.Errorf("Expected rate 0, got %d", limiter.rate)
	}

	if limiter.tokens != 0 {
		t.Errorf("Expected initial tokens 0, got %d", limiter.tokens)
	}
}

func TestNewWithNegativeRate(t *testing.T) {
	limiter := New(-5)
	if limiter == nil {
		t.Fatal("Expected limiter to be created with negative rate, got nil")
	}

	if limiter.rate != -5 {
		t.Errorf("Expected rate -5, got %d", limiter.rate)
	}

	if limiter.tokens != -5 {
		t.Errorf("Expected initial tokens -5, got %d", limiter.tokens)
	}
}

func TestAllow(t *testing.T) {
	rate := 2
	limiter := New(rate)

	// 前两个请求应该被允许
	if !limiter.allow() {
		t.Error("Expected first request to be allowed")
	}
	if !limiter.allow() {
		t.Error("Expected second request to be allowed")
	}

	// 第三个请求应该被拒绝
	if limiter.allow() {
		t.Error("Expected third request to be denied")
	}

	// 验证令牌数量
	if limiter.tokens != 0 {
		t.Errorf("Expected tokens to be 0, got %d", limiter.tokens)
	}
}

func TestAllowWithTokenReplenishment(t *testing.T) {
	rate := 1
	limiter := New(rate)

	// 第一个请求应该被允许
	if !limiter.allow() {
		t.Error("Expected first request to be allowed")
	}

	// 第二个请求应该被拒绝
	if limiter.allow() {
		t.Error("Expected second request to be denied")
	}

	// 等待一秒让令牌重新填充
	time.Sleep(1100 * time.Millisecond)

	// 现在应该被允许
	if !limiter.allow() {
		t.Error("Expected request after waiting to be allowed")
	}
}

func TestAllowWithPartialTokenReplenishment(t *testing.T) {
	rate := 10
	limiter := New(rate)

	// 使用所有令牌
	for i := 0; i < rate; i++ {
		if !limiter.allow() {
			t.Errorf("Expected request %d to be allowed", i+1)
		}
	}

	// 等待0.5秒，应该补充5个令牌
	time.Sleep(500 * time.Millisecond)

	// 应该允许5个请求
	for i := 0; i < 5; i++ {
		if !limiter.allow() {
			t.Errorf("Expected request after waiting to be allowed, got denied at %d", i+1)
		}
	}

	// 第6个请求应该被拒绝
	if limiter.allow() {
		t.Error("Expected 6th request to be denied")
	}
}

func TestMiddleware(t *testing.T) {
	rate := 2
	limiter := New(rate)

	// 创建测试处理器
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("OK"))
	})

	// 创建中间件
	middleware := limiter.Middleware(handler)

	// 创建测试请求
	req := httptest.NewRequest("GET", "/", nil)

	// 前两个请求应该成功
	for i := 0; i < 2; i++ {
		w := httptest.NewRecorder()
		middleware.ServeHTTP(w, req)

		if w.Code != http.StatusOK {
			t.Errorf("Expected status %d for request %d, got %d", http.StatusOK, i+1, w.Code)
		}

		if w.Body.String() != "OK" {
			t.Errorf("Expected body 'OK' for request %d, got '%s'", i+1, w.Body.String())
		}
	}

	// 第三个请求应该被限流
	w := httptest.NewRecorder()
	middleware.ServeHTTP(w, req)

	if w.Code != http.StatusTooManyRequests {
		t.Errorf("Expected status %d for rate limited request, got %d", http.StatusTooManyRequests, w.Code)
	}

	if strings.TrimSpace(w.Body.String()) != "Too Many Requests" {
		t.Errorf("Expected body 'Too Many Requests', got '%s'", w.Body.String())
	}
}

func TestMiddlewareWithZeroRate(t *testing.T) {
	limiter := New(0)

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("OK"))
	})

	middleware := limiter.Middleware(handler)
	req := httptest.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()

	middleware.ServeHTTP(w, req)

	// 所有请求都应该被限流
	if w.Code != http.StatusTooManyRequests {
		t.Errorf("Expected status %d for zero rate limiter, got %d", http.StatusTooManyRequests, w.Code)
	}
}

func TestConcurrentAllow(t *testing.T) {
	rate := 5
	limiter := New(rate)
	var wg sync.WaitGroup
	results := make(chan bool, 10)

	// 启动10个并发请求
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			results <- limiter.allow()
		}()
	}

	wg.Wait()
	close(results)

	// 统计允许的请求数
	allowed := 0
	for result := range results {
		if result {
			allowed++
		}
	}

	// 应该只允许rate个请求
	if allowed != rate {
		t.Errorf("Expected %d allowed requests, got %d", rate, allowed)
	}
}

func TestConcurrentMiddleware(t *testing.T) {
	rate := 3
	limiter := New(rate)

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("OK"))
	})

	middleware := limiter.Middleware(handler)
	req := httptest.NewRequest("GET", "/", nil)

	var wg sync.WaitGroup
	results := make(chan int, 10)

	// 启动10个并发请求
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			w := httptest.NewRecorder()
			middleware.ServeHTTP(w, req)
			results <- w.Code
		}()
	}

	wg.Wait()
	close(results)

	// 统计状态码
	okCount := 0
	tooManyRequestsCount := 0
	for code := range results {
		switch code {
		case http.StatusOK:
			okCount++
		case http.StatusTooManyRequests:
			tooManyRequestsCount++
		default:
			t.Errorf("Unexpected status code: %d", code)
		}
	}

	// 验证结果
	if okCount != rate {
		t.Errorf("Expected %d OK responses, got %d", rate, okCount)
	}

	expectedTooManyRequests := 10 - rate
	if tooManyRequestsCount != expectedTooManyRequests {
		t.Errorf("Expected %d Too Many Requests responses, got %d", expectedTooManyRequests, tooManyRequestsCount)
	}
}

func TestTokenReplenishmentAccuracy(t *testing.T) {
	testWithTimeout(t, 5*time.Second, func() {
		rate := 10
		limiter := New(rate)

		// 使用所有令牌
		for i := 0; i < rate; i++ {
			if !limiter.allow() {
				t.Fatalf("Expected to allow request %d", i)
			}
		}

		// 确保没有剩余令牌
		if limiter.allow() {
			t.Fatal("Expected no tokens left")
		}

		// 等待足够的时间让令牌补充
		time.Sleep(1100 * time.Millisecond) // 稍微多等一点时间

		// 触发令牌补充逻辑并验证
		allowedCount := 0
		for i := 0; i < rate+1; i++ { // 尝试比rate多1个请求
			if limiter.allow() {
				allowedCount++
			}
		}

		// 应该允许rate个请求（令牌已完全补充）
		if allowedCount != rate {
			t.Errorf("Expected %d allowed requests after 1+ second, got %d", rate, allowedCount)
		}
	})
}

func TestMinFunction(t *testing.T) {
	tests := []struct {
		a, b, expected int
	}{
		{1, 2, 1},
		{2, 1, 1},
		{5, 5, 5},
		{0, 10, 0},
		{-5, 5, -5},
		{5, -5, -5},
	}

	for _, test := range tests {
		result := min(test.a, test.b)
		if result != test.expected {
			t.Errorf("min(%d, %d) = %d, expected %d", test.a, test.b, result, test.expected)
		}
	}
}

func TestLimiterStateConsistency(t *testing.T) {
	testWithTimeout(t, 5*time.Second, func() {
		rate := 5
		limiter := New(rate)

		// 验证初始状态 - 应该允许rate个请求
		initialAllowed := 0
		for i := 0; i < rate; i++ {
			if limiter.allow() {
				initialAllowed++
			}
		}
		if initialAllowed != rate {
			t.Errorf("Expected %d initial allowed requests, got %d", rate, initialAllowed)
		}

		// 确保没有更多令牌
		if limiter.allow() {
			t.Error("Expected no more tokens available")
		}

		// 等待令牌重新填充
		time.Sleep(1200 * time.Millisecond) // 稍微多等一点时间

		// 验证令牌已重新填充 - 应该再次允许rate个请求
		refilledAllowed := 0
		for i := 0; i < rate; i++ {
			if limiter.allow() {
				refilledAllowed++
			}
		}
		if refilledAllowed != rate {
			t.Errorf("Expected %d refilled allowed requests, got %d", rate, refilledAllowed)
		}

		// 确保没有额外的令牌
		if limiter.allow() {
			t.Error("Expected no extra tokens after refill")
		}
	})
}

func TestHighRateLimiter(t *testing.T) {
	rate := 1000
	limiter := New(rate)

	// 测试高频率请求
	allowed := 0
	for i := 0; i < rate+10; i++ {
		if limiter.allow() {
			allowed++
		}
	}

	if allowed != rate {
		t.Errorf("Expected %d allowed requests, got %d", rate, allowed)
	}
}

func TestLimiterWithDifferentHTTPMethods(t *testing.T) {
	rate := 2
	limiter := New(rate)

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(r.Method))
	})

	middleware := limiter.Middleware(handler)

	methods := []string{"GET", "POST", "PUT", "DELETE", "PATCH"}

	// 测试不同HTTP方法
	for _, method := range methods {
		req := httptest.NewRequest(method, "/", nil)
		w := httptest.NewRecorder()
		middleware.ServeHTTP(w, req)

		// 前两个请求应该成功
		if w.Code == http.StatusOK {
			if w.Body.String() != method {
				t.Errorf("Expected method %s in response, got %s", method, w.Body.String())
			}
		} else if w.Code != http.StatusTooManyRequests {
			t.Errorf("Unexpected status code %d for method %s", w.Code, method)
		}
	}
}
